#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Spatial and temporal expansion of global wildland fire activity in response to climate change """

import numpy as np
import pandas as pd
from pandas import Series, DataFrame
import xarray as xr
import datetime as dt
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from mpl_toolkits.basemap import Basemap, maskoceans
from matplotlib.cm import get_cmap
import matplotlib.colors as mcolors

### File reading ###
####################

ruta_gfed4='../../DATA/GFED4/'	# GFED4 data path
ds_ba_years = xr.open_dataset(ruta_gfed4+'GFED4_BA_years.nc')
ds_ba_a = xr.open_dataset(ruta_gfed4+'GFED4_BA_a_mean.nc')
BA_years = ds_ba_years.BurnedArea		# Burned Area by year
BA_a = ds_ba_a.BurnedArea       # Burned Area annal mean

BA_max = BA_years.max(dim='year')
BA_max.values[BA_a.values==0.]=np.nan
BA_max.to_netcdf('BA_max.nc')

ruta_wfde5='../../DATA/WFDE5/'  # WFDE5 data path
ds_tp_months = xr.open_dataset(ruta_wfde5+'wfde5_tp_months_interp025.nc')
ds_t2m_months = xr.open_dataset(ruta_wfde5+'wfde5_t2m_months_interp025.nc')
P_months = ds_tp_months.tp                                      # Month total precipitation
T_months = ds_t2m_months.t2m                                    # Month mean temperature
P_years = P_months.groupby('time.year').sum(dim='time')         # Year total precipitation
T_years = T_months.groupby('time.year').mean(dim='time')        # Year mean temperature
Pmin_years = P_months.groupby('time.year').min(dim='time')      # Total precipitation of the driest month by year
Tmax_years = T_months.groupby('time.year').max(dim='time')      # Temperature of the hottest month by year

ds_fs = xr.open_dataset('FS_fullarray.nc')

FS = ds_fs.FS_array    

P_years = P_years.mean(dim='year',skipna=True)
P_years = xr.where(BA_max==0,np.nan,P_years)
P_months = xr.where((FS==0),np.nan,P_months)
P_months = P_months.groupby('time.year').mean(dim='time',skipna=True).mean(dim='year',skipna=True)
P_months = xr.where(BA_max==0,np.nan,P_months)
P_perce = P_months/P_years*100.
P_perce.to_netcdf('P_perce.nc')

T_years = T_years.mean(dim='year',skipna=True)
T_years = xr.where(BA_max==0,np.nan,T_years)
T_months = xr.where((FS==0),np.nan,T_months)
T_months = T_months.groupby('time.year').mean(dim='time',skipna=True).mean(dim='year',skipna=True)
T_months = xr.where(BA_max==0,np.nan,T_months)
T_anom = T_months-T_years
T_anom.to_netcdf('T_anom.nc')